// Copyright (C) 2025  lihanrui2913
#include "xhci-hcd.h"
#include <libs/aether/irq.h>

spinlock_t transfer_lock = {0};

// 获取端口速度名称
static const char *usb_speed_name(uint8_t speed) {
    switch (speed) {
    case USB_SPEED_LOW:
        return "Low Speed (1.5 Mbps)";
    case USB_SPEED_FULL:
        return "Full Speed (12 Mbps)";
    case USB_SPEED_HIGH:
        return "High Speed (480 Mbps)";
    case USB_SPEED_SUPER:
        return "Super Speed (5 Gbps)";
    default:
        return "Other Speed";
    }
}

static uint8_t usb_speed_to_xhci_port_speed(uint8_t speed) {
    switch (speed) {
    case USB_SPEED_FULL:
        return 1; // Full Speed (12 Mbps)
    case USB_SPEED_LOW:
        return 2; // Low Speed (1.5 Mbps)
    case USB_SPEED_HIGH:
        return 3; // High Speed (480 Mbps)
    case USB_SPEED_SUPER:
        return 4; // Super Speed (5 Gbps)
    default:
        return speed;
    }
}

// 从端口状态获取设备速度
static uint8_t xhci_port_speed_to_usb_speed(uint32_t portsc) {
    // XHCI PORTSC 的速度字段在 bits 10-13
    uint32_t port_speed = (portsc >> 10) & 0x0F;

    switch (port_speed) {
    case 1:
        return USB_SPEED_FULL; // Full Speed (12 Mbps)
    case 2:
        return USB_SPEED_LOW; // Low Speed (1.5 Mbps)
    case 3:
        return USB_SPEED_HIGH; // High Speed (480 Mbps)
    case 4:
        return USB_SPEED_SUPER; // Super Speed (5 Gbps)
    default:
        return port_speed;
    }
}

// 前向声明
static int xhci_hcd_init(usb_hcd_t *hcd);
static int xhci_hcd_shutdown(usb_hcd_t *hcd);
static int xhci_reset_port(usb_hcd_t *hcd, usb_hub_t *hub,
                           usb_device_t *device);
static int xhci_disconnect_port(usb_hcd_t *hcd, usb_hub_t *hub,
                                usb_device_t *device);
static int xhci_enable_slot(usb_hcd_t *hcd, usb_device_t *device);
static int xhci_disable_slot(usb_hcd_t *hcd, usb_device_t *device);
static int xhci_address_device(usb_hcd_t *hcd, usb_device_t *device,
                               uint8_t address);
static int xhci_configure_endpoint(usb_hcd_t *hcd, usb_endpoint_t *endpoint);
static int xhci_control_transfer(usb_hcd_t *hcd, usb_transfer_t *transfer,
                                 usb_device_request_t *setup);
static int xhci_bulk_transfer(usb_hcd_t *hcd, usb_transfer_t *transfer);
static int xhci_interrupt_transfer(usb_hcd_t *hcd, usb_transfer_t *transfer);

static usb_hcd_ops_t xhci_ops = {
    .init = xhci_hcd_init,
    .shutdown = xhci_hcd_shutdown,
    .enable_slot = xhci_enable_slot,
    .disable_slot = xhci_disable_slot,
    .address_device = xhci_address_device,
    .configure_endpoint = xhci_configure_endpoint,
    .control_transfer = xhci_control_transfer,
    .bulk_transfer = xhci_bulk_transfer,
    .interrupt_transfer = xhci_interrupt_transfer,
};

static usb_hub_ops_t xhci_root_hub_ops = {
    .reset_port = xhci_reset_port,
    .disconnect_port = xhci_disconnect_port, // TODO: 实现
};

// 分配DMA内存（对齐到64字节）
static void *xhci_alloc_dma(size_t size, uint64_t *phys_addr) {
    void *ptr = alloc_frames_bytes(size);
    if (ptr) {
        memset(ptr, 0, size);
        if (phys_addr) {
            *phys_addr =
                translate_address(get_current_page_dir(false), (uint64_t)ptr);
        }
    }
    return ptr;
}

// 分配传输环
xhci_ring_t *xhci_alloc_ring(uint32_t num_trbs) {
    xhci_ring_t *ring = (xhci_ring_t *)malloc(sizeof(xhci_ring_t));
    if (!ring) {
        return NULL;
    }

    memset(ring, 0, sizeof(xhci_ring_t));
    ring->size = num_trbs;
    ring->trbs = (xhci_trb_t *)xhci_alloc_dma(sizeof(xhci_trb_t) * num_trbs,
                                              &ring->phys_addr);
    if (!ring->trbs) {
        free(ring);
        return NULL;
    }

    ring->cycle_state = 1;
    ring->enqueue_index = 0;
    ring->dequeue_index = 0;

    return ring;
}

// 释放传输环
void xhci_free_ring(xhci_ring_t *ring) {
    if (!ring)
        return;

    if (ring->trbs) {
        free_frames_bytes(ring->trbs, ring->size * sizeof(xhci_trb_t));
    }
    free(ring);
}

// 向环中添加TRB
xhci_trb_t *xhci_queue_trb(xhci_ring_t *ring, uint64_t param, uint32_t status,
                           uint32_t control) {
    if (!ring || !ring->trbs) {
        printk("XHCI: Invalid ring in queue_trb\n");
        return NULL;
    }

    spin_lock(&transfer_lock);

    uint32_t idx = ring->enqueue_index;

    // 检查是否需要 Link TRB
    if (idx >= ring->size - 1) {
        // 设置 Link TRB
        xhci_trb_t *link = &ring->trbs[ring->size - 1];
        link->parameter = ring->phys_addr;
        link->status = 0;
        link->control = (TRB_TYPE_LINK << 10);

        link->control |= TRB_LINK_TOGGLE_CYCLE;

        // 设置 cycle bit
        if (ring->cycle_state) {
            link->control |= TRB_FLAG_CYCLE;
        }

        // 翻转 cycle state
        ring->cycle_state = !ring->cycle_state;
        ring->enqueue_index = 0;
        idx = 0;
    }

    xhci_trb_t *trb = &ring->trbs[idx];

    // 填充 TRB（先不设置 cycle bit）
    trb->parameter = param;
    trb->status = status;
    trb->control = control & ~TRB_FLAG_CYCLE;

    // 设置 cycle bit（激活 TRB）
    if (ring->cycle_state) {
        trb->control |= TRB_FLAG_CYCLE;
    }

    // 更新 enqueue 指针
    ring->enqueue_index = (idx + 1) % ring->size;

    spin_unlock(&transfer_lock);

    return trb;
}

// 敲门铃寄存器
void xhci_ring_doorbell(xhci_hcd_t *xhci, uint8_t slot_id, uint8_t dci) {
    xhci_writel(&xhci->doorbell_regs[slot_id], dci);
}

// 重置XHCI控制器
int xhci_reset(xhci_hcd_t *xhci) {
    uint32_t cmd, status;
    // 等待可以操作
    bool timeout = true;
    uint64_t time_ns = nanoTime() + 1ULL * 1000000000ULL;
    while (nanoTime() < time_ns) {
        cmd = xhci_readl(&xhci->op_regs->usbcmd);
        status = xhci_readl(&xhci->op_regs->usbsts);
        if (!(status & XHCI_STS_CNR)) {
            timeout = false;
            break;
        }
    }

    if (timeout) {
        printk("XHCI: Reset timeout\n");
        return -1;
    }

    // 停止
    cmd = xhci_readl(&xhci->op_regs->usbcmd);
    cmd &= ~XHCI_CMD_RUN;
    xhci_writel(&xhci->op_regs->usbcmd, cmd);

    // 等待停止
    timeout = true;
    time_ns = nanoTime() + 1ULL * 1000000000ULL;
    while (nanoTime() < time_ns) {
        cmd = xhci_readl(&xhci->op_regs->usbcmd);
        status = xhci_readl(&xhci->op_regs->usbsts);
        if (status & XHCI_STS_HCH) {
            timeout = false;
            break;
        }
    }

    if (timeout) {
        printk("XHCI: Reset timeout\n");
        return -1;
    }

    // 重置控制器
    cmd = xhci_readl(&xhci->op_regs->usbcmd);
    cmd |= XHCI_CMD_RESET;
    xhci_writel(&xhci->op_regs->usbcmd, cmd);

    // 等待重置完成
    timeout = true;
    time_ns = nanoTime() + 1ULL * 1000000000ULL;
    while (nanoTime() < time_ns) {
        cmd = xhci_readl(&xhci->op_regs->usbcmd);
        if (!(cmd & XHCI_CMD_RESET)) {
            timeout = false;
            break;
        }
    }

    if (timeout) {
        printk("XHCI: Reset timeout\n");
        return -1;
    }

    return 0;
}

// 启动XHCI控制器
int xhci_start(xhci_hcd_t *xhci) {
    uint32_t cmd;

    cmd = xhci_readl(&xhci->op_regs->usbcmd);
    cmd |= XHCI_CMD_RUN;
    xhci_writel(&xhci->op_regs->usbcmd, cmd);

    // 等待控制器运行
    while (1) {
        uint32_t status = xhci_readl(&xhci->op_regs->usbsts);
        if (!(status & XHCI_STS_HCH)) {
            return 0;
        }
    }

    printk("XHCI: Failed to start controller\n");
    return -1;
}

// 停止XHCI控制器
int xhci_stop(xhci_hcd_t *xhci) {
    uint32_t cmd;

    cmd = xhci_readl(&xhci->op_regs->usbcmd);
    cmd &= ~XHCI_CMD_RUN;
    xhci_writel(&xhci->op_regs->usbcmd, cmd);

    // 等待控制器停止
    uint64_t time_ns = nanoTime() + 1ULL * 1000000000ULL;

    while (nanoTime() < time_ns) {
        uint32_t status = xhci_readl(&xhci->op_regs->usbsts);
        if (status & XHCI_STS_HCH) {
            return 0;
        }
    }

    printk("XHCI: Failed to stop controller\n");
    return -1;
}

static void xhci_pair_ports(xhci_hcd_t *xhci) {
    // 首先将所有 paired_port 初始化为 0xFF（表示未配对）
    for (int i = 0; i < xhci->max_ports; i++) {
        xhci->port_info[i].paired_port = 0xFF;
    }

    // 寻找配对的端口
    for (int i = 0; i < xhci->max_ports; i++) {
        if (xhci->port_info[i].protocol == XHCI_PROTOCOL_USB2) {
            // 对于 USB2.0 端口，寻找对应的 USB3.0 端口
            for (int j = 0; j < xhci->max_ports; j++) {
                if (xhci->port_info[j].protocol == XHCI_PROTOCOL_USB3) {
                    // 检查是否是配对的端口
                    // 通常配对的端口有相同的物理位置索引
                    // 计算在各自协议中的相对位置
                    uint8_t usb2_relative_pos =
                        i - (xhci->port_info[i].port_offset - 1);
                    uint8_t usb3_relative_pos =
                        j - (xhci->port_info[j].port_offset - 1);

                    if (usb2_relative_pos == usb3_relative_pos) {
                        xhci->port_info[i].paired_port = j;
                        xhci->port_info[j].paired_port = i;

                        printk("XHCI: Paired ports - USB2.0 port %d <-> USB3.0 "
                               "port %d\n",
                               i + 1, j + 1);

                        break;
                    }
                }
            }
        }
    }
}

// 扩展能力 ID
#define XHCI_EXT_CAPS_LEGACY 1
#define XHCI_EXT_CAPS_PROTOCOL 2

// 解析 Supported Protocol Capability
static int xhci_parse_protocol_caps(xhci_hcd_t *xhci) {
    uint32_t hccparams1 = xhci_readl(&xhci->cap_regs->hccparams1);
    uint32_t ext_offset = (hccparams1 >> 16) & 0xFFFF;

    if (ext_offset == 0) {
        printk("XHCI: No extended capabilities\n");
        return 0;
    }

    // 扩展能力在能力寄存器之后
    uint32_t *ext_cap = (uint32_t *)xhci->cap_regs + ext_offset;

    while (ext_offset) {
        uint32_t cap_header = xhci_readl(ext_cap);
        uint32_t cap_id = cap_header & 0xFF;
        uint32_t next_offset = (cap_header >> 8) & 0xFF;

        if (cap_id == XHCI_EXT_CAPS_LEGACY) {
            uint32_t val = xhci_readl(ext_cap);

            if (val & (1 << 16)) {
                printk("XHCI is BIOS owned!!!\n");
                xhci_writel(ext_cap, val | (1 << 24));

                bool timeout = true;
                uint64_t timeout_ns = nanoTime() + 5ULL * 1000000000ULL;
                while (nanoTime() < timeout_ns) {
                    val = xhci_readl(ext_cap);
                    if (!(val & (1 << 16))) {
                        timeout = false;
                        break;
                    }
                }

                if (timeout) {
                    printk("Failed to hand over XHCI from BIOS!!!\n");
                    xhci_writel(ext_cap, val & ~(1 << 16));
                }
            }

            val = xhci_readl(ext_cap + 1);
            val &= ((0x7 << 1) + (0xff << 5) + (0x7 << 17));
            val |= (0x7 << 29);
            xhci_writel(ext_cap + 1, val);
        } else if (cap_id == XHCI_EXT_CAPS_PROTOCOL) {
            // Supported Protocol Capability
            uint32_t cap_word0 = xhci_readl(ext_cap);
            uint32_t cap_word1 = xhci_readl(ext_cap + 1);
            uint32_t cap_word2 = xhci_readl(ext_cap + 2);
            uint32_t cap_word3 = xhci_readl(ext_cap + 3);

            // 解析协议版本
            uint8_t minor = (cap_word0 >> 16) & 0xFF;
            uint8_t major = (cap_word0 >> 24) & 0xFF;

            // 解析端口范围
            uint8_t port_offset = cap_word2 & 0xFF;
            uint8_t port_count = (cap_word2 >> 8) & 0xFF;

            uint8_t slot_type = cap_word3 & 0x1F;

            printk("XHCI: Protocol Capability - Major: %d, Minor: %d, Ports: "
                   "%d-%d, Slot Type: %d\n",
                   major, minor, port_offset, port_offset + port_count - 1,
                   slot_type);

            // 标记端口协议
            for (int i = 0;
                 i < port_count && (port_offset + i - 1) < xhci->max_ports;
                 i++) {
                uint8_t port_idx = port_offset + i - 1; // 转换为 0-based

                xhci->port_info[port_idx].port_num = port_idx;
                xhci->port_info[port_idx].protocol = major;
                xhci->port_info[port_idx].port_offset = port_offset;
                xhci->port_info[port_idx].port_count = port_count;
                xhci->port_info[port_idx].slot_type = slot_type;
            }
        }

        if (next_offset == 0) {
            break;
        }

        ext_cap += next_offset;
        ext_offset = next_offset;
    }

    xhci_pair_ports(xhci);

    return 0;
}

void xhci_handle_port_status(xhci_hcd_t *xhci, uint8_t port_id);

// 初始化XHCI HCD
static int xhci_hcd_init(usb_hcd_t *hcd) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;

    xhci->hcd = hcd;

    // 重置控制器
    if (xhci_reset(xhci) != 0) {
        return -1;
    }

    xhci->tracker_mutex.lock = 0;

    // 读取能力参数
    uint32_t hcsparams1 = xhci_readl(&xhci->cap_regs->hcsparams1);
    uint32_t hcsparams2 = xhci_readl(&xhci->cap_regs->hcsparams2);
    uint32_t hccparams1 = xhci_readl(&xhci->cap_regs->hccparams1);

    xhci->use_64byte_context = (hccparams1 & 0x04) != 0;
    printk("XHCI: context size: %s\n", xhci->use_64byte_context ? "64" : "32");

    xhci->max_slots = (hcsparams1 >> 0) & 0xFF;
    xhci->max_intrs = (hcsparams1 >> 8) & 0x7FF;
    xhci->max_ports = (hcsparams1 >> 24) & 0xFF;

    // 计算scratchpad数量
    uint32_t max_scratch_hi = (hcsparams2 >> 21) & 0x1F;
    uint32_t max_scratch_lo = (hcsparams2 >> 27) & 0x1F;
    xhci->num_scratchpads = (max_scratch_hi << 5) | max_scratch_lo;

    // 配置max slots
    uint32_t config = xhci->max_slots;
    config |=
        (xhci_readl(&xhci->cap_regs->hccparams2) & (1 << 5)) ? (1 << 9) : 0;
    xhci_writel(&xhci->op_regs->config, config);

    // 分配设备上下文基地址数组 (DCBAA)
    xhci->dcbaa = (uint64_t *)xhci_alloc_dma(
        (xhci->max_slots + 1) * sizeof(uint64_t), &xhci->dcbaa_phys);
    if (!xhci->dcbaa) {
        printk("XHCI: Failed to allocate DCBAA\n");
        return -1;
    }

    // 分配scratchpad buffers
    if (xhci->num_scratchpads > 0) {
        xhci->scratchpad_array = (uint64_t *)xhci_alloc_dma(
            xhci->num_scratchpads * sizeof(uint64_t), NULL);
        xhci->scratchpad_buffers =
            (void **)alloc_frames_bytes(xhci->num_scratchpads * sizeof(void *));

        for (uint32_t i = 0; i < xhci->num_scratchpads; i++) {
            uint64_t phys;
            xhci->scratchpad_buffers[i] = xhci_alloc_dma(4096, &phys);
            xhci->scratchpad_array[i] = phys;
        }

        xhci->dcbaa[0] = translate_address(get_current_page_dir(false),
                                           (uint64_t)xhci->scratchpad_array);
    }

    // 设置DCBAA指针
    xhci_writeq(&xhci->op_regs->dcbaap, xhci->dcbaa_phys);

    // 分配命令环
    xhci->cmd_ring = xhci_alloc_ring(256);
    if (!xhci->cmd_ring) {
        printk("XHCI: Failed to allocate command ring\n");
        return -1;
    }

    // 设置命令环指针
    xhci_writeq(&xhci->op_regs->crcr, xhci->cmd_ring->phys_addr | 1);

    // 分配事件环
    xhci->event_ring = xhci_alloc_ring(256);
    if (!xhci->event_ring) {
        printk("XHCI: Failed to allocate event ring\n");
        return -1;
    }

    // 分配事件环段表 (ERST)
    xhci->erst = (xhci_erst_entry_t *)xhci_alloc_dma(sizeof(xhci_erst_entry_t),
                                                     &xhci->erst_phys);
    if (!xhci->erst) {
        printk("XHCI: Failed to allocate ERST\n");
        return -1;
    }

    xhci->erst[0].address = xhci->event_ring->phys_addr;
    xhci->erst[0].size = xhci->event_ring->size;
    xhci->erst[0].reserved = 0;

    // 配置中断器0
    xhci_writel(&xhci->intr_regs[0].erstsz, 1);
    xhci_writeq(&xhci->intr_regs[0].erstba, xhci->erst_phys);
    xhci_writeq(&xhci->intr_regs[0].erdp, xhci->event_ring->phys_addr);

    xhci_writel(&xhci->intr_regs[0].imod, 0);

    // // 启用中断
    // uint32_t iman = xhci_readl(&xhci->intr_regs[0].iman);
    // iman |= (1 << 1);
    // xhci_writel(&xhci->intr_regs[0].iman, iman);

    // 分配端口信息数组
    xhci->port_info =
        (xhci_port_info_t *)malloc(sizeof(xhci_port_info_t) * xhci->max_ports);
    if (!xhci->port_info) {
        printk("XHCI: Failed to allocate port info\n");
        return -1;
    }
    memset(xhci->port_info, 0, sizeof(xhci_port_info_t) * xhci->max_ports);

    // 解析端口协议
    xhci_parse_protocol_caps(xhci);

    xhci->root_hub = malloc(sizeof(struct usb_hub));
    memset(xhci->root_hub, 0, sizeof(struct usb_hub));
    xhci_hub_private_t *hub_priv = malloc(sizeof(xhci_hub_private_t));
    hub_priv->hub = xhci->root_hub;
    hub_priv->xhci = xhci;
    xhci->root_hub->hcd_private = hub_priv;
    xhci->root_hub->device = NULL;
    xhci->root_hub->parent = NULL;
    xhci->root_hub->ops = &xhci_root_hub_ops;

    // uint32_t cmd = xhci_readl(&xhci->op_regs->usbcmd);
    // cmd |= XHCI_CMD_INTE;
    // xhci_writel(&xhci->op_regs->usbcmd, cmd);

    // 启动控制器
    if (xhci_start(xhci) != 0) {
        return -1;
    }

    // 启动事件处理线程
    if (xhci_start_event_handler(xhci) != 0) {
        printk("XHCI: Failed to start event handler\n");
        return -1;
    }

    // 先初始化USB3端口
    for (uint8_t p = 0; p < xhci->max_ports; p++) {
        xhci_port_info_t *port_info = &xhci->port_info[p];
        if (port_info->protocol == XHCI_PROTOCOL_USB3 &&
            (xhci_readl(&xhci->port_regs[p].portsc) & XHCI_PORTSC_CCS)) {
            xhci_handle_port_status(xhci, p);
        }
    }

    // 再初始化USB2端口
    for (uint8_t p = 0; p < xhci->max_ports; p++) {
        xhci_port_info_t *port_info = &xhci->port_info[p];
        if (port_info->protocol != XHCI_PROTOCOL_USB3 &&
            (xhci_readl(&xhci->port_regs[p].portsc) & XHCI_PORTSC_CCS) &&
            (port_info->paired_port == 0xFF ||
             !xhci->connection[port_info->paired_port])) {
            xhci_handle_port_status(xhci, p);
        }
    }

    return 0;
}

// 关闭XHCI HCD
static int xhci_hcd_shutdown(usb_hcd_t *hcd) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;

    // 停止事件处理线程
    xhci_stop_event_handler(xhci);

    xhci_stop(xhci);

    // // 释放资源
    // if (xhci->cmd_ring) {
    //     xhci_free_ring(xhci->cmd_ring);
    // }

    // if (xhci->event_ring) {
    //     xhci_free_ring(xhci->event_ring);
    // }

    // if (xhci->erst) {
    //     free(xhci->erst);
    // }

    // if (xhci->scratchpad_buffers) {
    //     for (uint32_t i = 0; i < xhci->num_scratchpads; i++) {
    //         if (xhci->scratchpad_buffers[i]) {
    //             free(xhci->scratchpad_buffers[i]);
    //         }
    //     }
    //     free(xhci->scratchpad_buffers);
    // }

    // if (xhci->scratchpad_array) {
    //     free(xhci->scratchpad_array);
    // }

    // if (xhci->dcbaa) {
    //     free(xhci->dcbaa);
    // }

    free(xhci);

    return 0;
}

static void delay(uint64_t ms) {
    uint64_t ns = ms * 1000000;
    uint64_t start = nanoTime();
    while (nanoTime() - start < ns) {
        arch_yield();
    }
}

// USB 2.0 端口重置
static int xhci_reset_port_usb2(xhci_hcd_t *xhci, uint8_t port,
                                usb_device_t *device) {
    uint32_t portsc;

    printk("XHCI: USB2 port reset begin\n");

    // 1. 清除所有状态变化位
    portsc = xhci_readl(&xhci->port_regs[port].portsc);
    portsc |= XHCI_PORTSC_CHANGE_BITS;
    portsc &= ~(XHCI_PORTSC_PED | XHCI_PORTSC_PR);
    xhci_writel(&xhci->port_regs[port].portsc, portsc);

    delay(1); // 让硬件处理

    // 2. 发起重置
    portsc = xhci_readl(&xhci->port_regs[port].portsc);

    // 保留所有非变化位，清除状态变化位（写0），设置 PR
    portsc &= ~XHCI_PORTSC_CHANGE_BITS; // 写0保留状态变化位
    portsc |= XHCI_PORTSC_PR;           // 设置重置位

    xhci_writel(&xhci->port_regs[port].portsc, portsc);

    printk("XHCI: Reset initiated, PORTSC=0x%08x\n", portsc);

    // 3. 等待 PR 位清零
    uint64_t timeout = nanoTime() + 100ULL * 1000000ULL; // 100ms 足够
    bool pr_cleared = false;

    while (nanoTime() < timeout) {
        portsc = xhci_readl(&xhci->port_regs[port].portsc);

        if (!(portsc & XHCI_PORTSC_PR)) {
            pr_cleared = true;
            printk("XHCI: PR cleared, PORTSC=0x%08x\n", portsc);
            break;
        }

        delay(1);
    }

    if (!pr_cleared) {
        printk("XHCI: Reset timeout (PR not cleared)\n");
        return -1;
    }

    // 4. 等待 PRC（重置完成）
    timeout = nanoTime() + 100ULL * 1000000ULL;
    bool prc_set = false;

    while (nanoTime() < timeout) {
        portsc = xhci_readl(&xhci->port_regs[port].portsc);

        if (portsc & XHCI_PORTSC_PRC) {
            prc_set = true;
            printk("XHCI: PRC set, PORTSC=0x%08x\n", portsc);
            break;
        }

        delay(1);
    }

    if (!prc_set) {
        printk("XHCI: PRC not set\n");
        return -1;
    }

    // 5. 等待端口使能（PED）和链路进入 U0
    timeout = nanoTime() + 100ULL * 1000000ULL; // 100ms
    bool enabled = false;

    while (nanoTime() < timeout) {
        portsc = xhci_readl(&xhci->port_regs[port].portsc);

        // 检查连接状态
        if (!(portsc & XHCI_PORTSC_CCS)) {
            printk("XHCI: Device disconnected\n");
            return -1;
        }

        // 检查 PED 和 PLS
        uint8_t pls = (portsc >> 5) & 0xF;

        printk("XHCI: [wait] PORTSC=0x%08x PLS=%d PED=%d\n", portsc, pls,
               !!(portsc & XHCI_PORTSC_PED));

        if ((portsc & XHCI_PORTSC_PED) && (pls == 0)) {
            enabled = true;
            printk("XHCI: Port enabled! PLS=U0\n");
            break;
        }

        delay(10);
    }

    if (!enabled) {
        printk("XHCI: Port not enabled after reset\n");
        printk("      Final PORTSC=0x%08x\n", portsc);

        uint8_t pls = (portsc >> 5) & 0xF;
        printk("      PLS=%d (should be 0 for U0)\n", pls);
        printk("      PED=%d (should be 1)\n", !!(portsc & XHCI_PORTSC_PED));
        printk("      CCS=%d\n", !!(portsc & XHCI_PORTSC_CCS));

        return -1;
    }

    // 6. 现在才清除状态变化位
    portsc = xhci_readl(&xhci->port_regs[port].portsc);
    portsc &= ~XHCI_PORTSC_CHANGE_BITS; // 保留状态位
    portsc |= XHCI_PORTSC_CHANGE_BITS;  // 设置变化位（写1清除）
    xhci_writel(&xhci->port_regs[port].portsc, portsc);

    printk("XHCI: Status bits cleared\n");

    // 7. USB规范要求重置后恢复时间
    delay(10);

    // 8. 读取设备速度
    portsc = xhci_readl(&xhci->port_regs[port].portsc);
    uint8_t speed = xhci_port_speed_to_usb_speed(portsc);
    device->speed = speed;

    return 0;
}

// USB 3.0 端口重置
static int xhci_reset_port_usb3(xhci_hcd_t *xhci, uint8_t port,
                                usb_device_t *device) {
    uint32_t portsc = xhci_readl(&xhci->port_regs[port].portsc);

    while (1) {
        portsc = xhci_readl(&xhci->port_regs[port].portsc);
        if (portsc & XHCI_PORTSC_CCS) {
            break;
        }
        printk("XHCI: Waiting for device connection on USB3 port...\n");
        delay(1000);
    }

    uint32_t pls = (portsc >> 5) & 0xF;
    bool disconnected = false;

    printk("XHCI: USB3 port in Link State %d\n", pls);

    // 1. 如果已在U0且已使能，无需重置
    if (pls == 0 && (portsc & XHCI_PORTSC_PED) && (portsc & XHCI_PORTSC_CCS)) {
        printk("XHCI: Already in U0 and enabled\n");
        device->speed = xhci_port_speed_to_usb_speed(portsc);
        return 0;
    }

    // 2. 从低功耗状态唤醒
    if (pls >= 1 && pls <= 3) { // U1/U2/U3
        printk("XHCI: Waking from U%d to U0\n", pls);

        xhci_writel(&xhci->port_regs[port].portsc,
                    XHCI_PORTSC_PP | (1 << 16) | (0 << 5));

        uint64_t timeout = nanoTime() + 500ULL * 1000000ULL;
        while (nanoTime() < timeout) {
            portsc = xhci_readl(&xhci->port_regs[port].portsc);
            pls = (portsc >> 5) & 0xF;
            if (pls == 0 && (portsc & XHCI_PORTSC_PED)) {
                device->speed = xhci_port_speed_to_usb_speed(portsc);
                printk("XHCI: Wake complete\n");
                return 0;
            }
            delay(1);
        }

        // 唤醒失败，继续尝试Warm Reset
        printk("XHCI: Wake failed, trying Warm Reset\n");
    }

    // 3. Warm Reset（处理Disabled/RxDetect/Compliance等状态）
    printk("XHCI: Performing Warm Reset (PLS=%d)\n", pls);

    // 清除状态位
    xhci_writel(&xhci->port_regs[port].portsc,
                XHCI_PORTSC_CSC | XHCI_PORTSC_PEC | XHCI_PORTSC_WRC |
                    XHCI_PORTSC_OCC | XHCI_PORTSC_PRC);

    // 发起Warm Reset
    xhci_writel(&xhci->port_regs[port].portsc, XHCI_PORTSC_PP | (1 << 31));

    // 等待WPR清零
    uint64_t timeout = nanoTime() + 500ULL * 1000000ULL;
    while (nanoTime() < timeout) {
        portsc = xhci_readl(&xhci->port_regs[port].portsc);

        // 监测断开（不立即失败）
        if (!(portsc & XHCI_PORTSC_CCS) && !disconnected) {
            printk("XHCI: Device briefly disconnected during Warm Reset\n");
            disconnected = true;
        }

        if (!(portsc & (1 << 31))) {
            printk("XHCI: Warm Reset signal complete\n");
            break;
        }
        delay(1);
    }

    // 等待WRC置位或设备重新连接
    timeout = nanoTime() + 1000ULL * 1000000ULL;
    bool wrc_set = false;

    while (nanoTime() < timeout) {
        portsc = xhci_readl(&xhci->port_regs[port].portsc);

        if (portsc & XHCI_PORTSC_WRC) {
            wrc_set = true;
            printk("XHCI: WRC set\n");
            break;
        }

        // 如果之前断开，检查是否重新连接
        if (disconnected && (portsc & XHCI_PORTSC_CCS)) {
            printk("XHCI: Device reconnected during Warm Reset\n");
            delay(50); // 让连接稳定
        }

        delay(10);
    }

    if (wrc_set) {
        // 清除WRC
        xhci_writel(&xhci->port_regs[port].portsc, XHCI_PORTSC_WRC);
    }

    // 4. 如果检测到断开，等待稳定重连
    portsc = xhci_readl(&xhci->port_regs[port].portsc);
    if (!(portsc & XHCI_PORTSC_CCS)) {
        printk("XHCI: Waiting for USB3 device reconnection...\n");

        timeout = nanoTime() + 2000ULL * 1000000ULL; // USB3需要更长时间
        bool reconnected = false;

        while (nanoTime() < timeout) {
            portsc = xhci_readl(&xhci->port_regs[port].portsc);

            if (portsc & XHCI_PORTSC_CCS) {
                delay(100); // 等待连接稳定

                portsc = xhci_readl(&xhci->port_regs[port].portsc);
                if (portsc & XHCI_PORTSC_CCS) {
                    reconnected = true;
                    printk("XHCI: USB3 device reconnected\n");
                    break;
                }
            }

            delay(10);
        }

        if (!reconnected) {
            printk("XHCI: USB3 device did not reconnect\n");
            return -1;
        }
    }

    // 5. 等待链路训练到U0
    printk("XHCI: Waiting for link training to U0...\n");
    timeout = nanoTime() + 2000ULL * 1000000ULL; // 2秒

    while (nanoTime() < timeout) {
        portsc = xhci_readl(&xhci->port_regs[port].portsc);
        pls = (portsc >> 5) & 0xF;

        // 检查持续断开
        if (!(portsc & XHCI_PORTSC_CCS)) {
            printk("XHCI: Device disconnected during link training\n");
            return -1;
        }

        // 到达U0且已使能
        if (pls == 0 && (portsc & XHCI_PORTSC_PED)) {
            printk("XHCI: Link training complete, in U0\n");
            break;
        }

        // 打印链路状态变化
        static uint32_t last_pls = 0xFF;
        if (pls != last_pls) {
            const char *pls_names[] = {"U0",
                                       "U1",
                                       "U2",
                                       "U3",
                                       "Disabled",
                                       "RxDetect",
                                       "Inactive",
                                       "Polling",
                                       "Recovery",
                                       "HotReset",
                                       "ComplianceMode",
                                       "TestMode",
                                       "Reserved",
                                       "Reserved",
                                       "Reserved",
                                       "ResumeState"};
            printk("XHCI:   Link State -> %s\n", pls_names[pls]);
            last_pls = pls;
        }

        delay(10);
    }

    // 6. 最终检查
    portsc = xhci_readl(&xhci->port_regs[port].portsc);
    pls = (portsc >> 5) & 0xF;

    printk("XHCI: Final PORTSC=0x%08x, PLS=%d\n", portsc, pls);

    if (!(portsc & XHCI_PORTSC_CCS)) {
        printk("XHCI: Device not connected\n");
        return -1;
    }

    if (!(portsc & XHCI_PORTSC_PED)) {
        printk("XHCI: Port not enabled\n");
        return -1;
    }

    if (pls != 0) {
        printk("XHCI: Not in U0 (still in PLS=%d)\n", pls);
        return -1;
    }

    // 7. 清除所有状态变化位
    xhci_writel(&xhci->port_regs[port].portsc,
                XHCI_PORTSC_CSC | XHCI_PORTSC_PEC | XHCI_PORTSC_WRC |
                    XHCI_PORTSC_OCC | XHCI_PORTSC_PRC);

    device->speed = xhci_port_speed_to_usb_speed(portsc);
    printk("XHCI: USB3 reset complete, speed=%s\n",
           usb_speed_name(device->speed));

    return 0;
}

// 主重置函数（添加重试机制）
static int xhci_reset_port(usb_hcd_t *hcd, usb_hub_t *hub,
                           usb_device_t *device) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;
    uint8_t port = device->port;

    if (port >= xhci->max_ports) {
        printk("XHCI: Invalid port %d\n", port);
        return -1;
    }

    xhci_port_info_t *port_info = &xhci->port_info[port];
    bool is_usb3 = (port_info->protocol == XHCI_PROTOCOL_USB3);

    printk("XHCI: Resetting port %d (%s)\n", port, is_usb3 ? "USB3" : "USB2");

    // 1. 检查设备连接
    uint32_t portsc = xhci_readl(&xhci->port_regs[port].portsc);
    printk("XHCI: Initial PORTSC=0x%08x\n", portsc);

    // 2. 检查端口供电
    if (!(portsc & XHCI_PORTSC_PP)) {
        printk("XHCI: Enabling port power\n");
        xhci_writel(&xhci->port_regs[port].portsc, XHCI_PORTSC_PP);
        delay(20);
    }

    // 3. 重置（最多重试2次）
    int ret = -1;
    for (int retry = 0; retry < 2; retry++) {
        if (retry > 0) {
            printk("XHCI: Retry %d...\n", retry);
            delay(500); // 重试前等待
        }

        if (is_usb3) {
            ret = xhci_reset_port_usb3(xhci, port, device);
        } else {
            ret = xhci_reset_port_usb2(xhci, port, device);
        }

        if (ret == 0) {
            return 0;
        }
    }

    printk("XHCI: Port reset failed after retries\n");
    return -1;
}

static int xhci_disconnect_port(usb_hcd_t *hcd, usb_hub_t *hub,
                                usb_device_t *device) {
    return -1; // TODO: 实现
}

// 为设备初始化EP0（在enable_slot之后调用）
int xhci_setup_default_endpoint(usb_hcd_t *hcd, usb_device_t *device) {
    // 为EP0分配端点私有数据
    xhci_endpoint_private_t *ep0_priv =
        (xhci_endpoint_private_t *)malloc(sizeof(xhci_endpoint_private_t));
    if (!ep0_priv) {
        printk("XHCI: Failed to allocate EP0 private data\n");
        return -1;
    }

    memset(ep0_priv, 0, sizeof(xhci_endpoint_private_t));
    ep0_priv->dci = 1; // EP0的DCI总是1

    // 分配传输环
    ep0_priv->transfer_ring = xhci_alloc_ring(256);
    if (!ep0_priv->transfer_ring) {
        printk("XHCI: Failed to allocate EP0 transfer ring\n");
        free(ep0_priv);
        return -1;
    }

    // 设置EP0基本信息
    device->endpoints[0].address = 0;
    device->endpoints[0].attributes = USB_ENDPOINT_XFER_CONTROL;
    device->endpoints[0].max_packet_size = 8; // 默认最小值，后续会更新
    device->endpoints[0].interval = 0;
    device->endpoints[0].device = device;
    device->endpoints[0].hcd_private = ep0_priv;

    return 0;
}

spinlock_t xhci_command_lock = {0};

// 启用槽位 - 使用等待机制
static int xhci_enable_slot(usb_hcd_t *hcd, usb_device_t *device) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;

    // 分配设备私有数据
    xhci_device_private_t *dev_priv =
        (xhci_device_private_t *)malloc(sizeof(xhci_device_private_t));
    if (!dev_priv) {
        return -1;
    }
    memset(dev_priv, 0, sizeof(xhci_device_private_t));

    // 分配设备上下文
    dev_priv->device_ctx = xhci_alloc_dma(
        sizeof(xhci_device_ctx_t) * (xhci->use_64byte_context ? 2 : 1),
        &dev_priv->device_ctx_phys);
    if (!dev_priv->device_ctx) {
        free(dev_priv);
        return -1;
    }

    // 分配输入上下文
    dev_priv->input_ctx = xhci_alloc_dma(sizeof(xhci_input_ctx_t) *
                                             (xhci->use_64byte_context ? 2 : 1),
                                         &dev_priv->input_ctx_phys);
    if (!dev_priv->input_ctx) {
        free_frames_bytes(dev_priv->device_ctx,
                          sizeof(xhci_device_ctx_t) *
                              (xhci->use_64byte_context ? 2 : 1));
        free(dev_priv);
        return -1;
    }

    // 分配完成结构
    xhci_command_completion_t *completion = xhci_alloc_command_completion();
    if (!completion) {
        free_frames_bytes(dev_priv->input_ctx,
                          sizeof(xhci_input_ctx_t) *
                              (xhci->use_64byte_context ? 2 : 1));
        free_frames_bytes(dev_priv->device_ctx,
                          sizeof(xhci_device_ctx_t) *
                              (xhci->use_64byte_context ? 2 : 1));
        free(dev_priv);
        return -1;
    }

    completion->device = device;

    spin_lock(&xhci_command_lock);

    xhci_port_info_t *port_info = &xhci->port_info[device->port];

    // 发送Enable Slot命令
    xhci_trb_t *cmd_trb =
        xhci_queue_trb(xhci->cmd_ring, 0, 0,
                       ((uint32_t)port_info->slot_type << 16) |
                           (TRB_TYPE_ENABLE_SLOT << 10) | TRB_FLAG_IOC);

    // 跟踪命令
    xhci_track_command(xhci, cmd_trb, completion, TRB_TYPE_ENABLE_SLOT);

    // 敲门铃
    xhci_ring_doorbell(xhci, 0, 0);

    spin_unlock(&xhci_command_lock);

    // 等待命令完成（5秒超时）
    int ret = xhci_wait_for_command(completion, 5000);

    if (ret == 0) {
        // 命令成功，获取slot ID
        dev_priv->slot_id = completion->slot_id;

        // 更新DCBAA
        xhci->dcbaa[dev_priv->slot_id] = dev_priv->device_ctx_phys;

        device->hcd_private = dev_priv;

        // 初始化EP0
        ret = xhci_setup_default_endpoint(hcd, device);
        if (ret != 0) {
            printk("XHCI: Failed to setup EP0\n");
            xhci->dcbaa[dev_priv->slot_id] = 0;
            free_frames_bytes(dev_priv->input_ctx, sizeof(xhci_input_ctx_t));
            free_frames_bytes(dev_priv->device_ctx, sizeof(xhci_device_ctx_t));
            dev_priv->input_ctx = NULL;
            dev_priv->device_ctx = NULL;
            free(dev_priv);
            device->hcd_private = NULL;
            xhci_free_command_completion(completion);
            return -1;
        }
    } else {
        // 命令失败
        free_frames_bytes(dev_priv->input_ctx, sizeof(xhci_input_ctx_t));
        free_frames_bytes(dev_priv->device_ctx, sizeof(xhci_device_ctx_t));
        free(dev_priv);
    }

    xhci_free_command_completion(completion);
    return ret;
}

// 禁用槽位
static int xhci_disable_slot(usb_hcd_t *hcd, usb_device_t *device) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;
    xhci_device_private_t *dev_priv =
        (xhci_device_private_t *)device->hcd_private;

    if (!dev_priv) {
        return -1;
    }

    printk("XHCI: Disabling slot %d\n", dev_priv->slot_id);

    spin_lock(&xhci_command_lock);

    // 发送Disable Slot命令
    xhci_trb_t *trb =
        xhci_queue_trb(xhci->cmd_ring, 0, 0,
                       (TRB_TYPE_DISABLE_SLOT << 10) |
                           ((uint32_t)dev_priv->slot_id << 24) | TRB_FLAG_IOC);

    xhci_ring_doorbell(xhci, 0, 0);

    spin_unlock(&xhci_command_lock);

    xhci_handle_events(xhci);

    // 清除DCBAA
    xhci->dcbaa[dev_priv->slot_id] = 0;

    // 释放资源
    if (dev_priv->device_ctx) {
        free_frames_bytes(dev_priv->device_ctx,
                          sizeof(xhci_device_ctx_t) *
                              (xhci->use_64byte_context ? 2 : 1));
    }
    if (dev_priv->input_ctx) {
        free_frames_bytes(dev_priv->input_ctx,
                          sizeof(xhci_input_ctx_t) *
                              (xhci->use_64byte_context ? 2 : 1));
    }
    free(dev_priv);

    device->hcd_private = NULL;

    return 0;
}

// 地址设备 - 这是设备枚举的关键步骤
static int xhci_address_device(usb_hcd_t *hcd, usb_device_t *device,
                               uint8_t address) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;
    xhci_device_private_t *dev_priv =
        (xhci_device_private_t *)device->hcd_private;

    if (!dev_priv) {
        printk("XHCI: ERROR - Device private data is NULL\n");
        return -1;
    }

    // 验证EP0已经初始化
    if (!device->endpoints[0].hcd_private) {
        printk("XHCI: CRITICAL - EP0 not initialized!\n");
        return -1;
    }

    xhci_endpoint_private_t *ep0_priv =
        (xhci_endpoint_private_t *)device->endpoints[0].hcd_private;

    if (!ep0_priv->transfer_ring) {
        printk("XHCI: CRITICAL - EP0 transfer ring not allocated!\n");
        return -1;
    }

    if (xhci->use_64byte_context) {
        // 清零输入上下文
        memset(dev_priv->input_ctx_64, 0, sizeof(xhci_input_ctx_64_t));

        // ===== 1. 设置输入控制上下文 =====
        dev_priv->input_ctx_64->ctrl.drop_flags = 0;
        dev_priv->input_ctx_64->ctrl.add_flags = 0x03; // A0 (Slot) | A1 (EP0)

        // ===== 2. 配置 Slot 上下文 =====
        uint32_t route_string = 0; // Root hub port
        uint32_t speed = usb_speed_to_xhci_port_speed(device->speed);
        uint32_t context_entries = 1; // EP0的DCI是1

        // dev_info: Route String, Speed, Context Entries
        dev_priv->input_ctx_64->slot.dev_info =
            SLOT_CTX_ROUTE_STRING(route_string) | SLOT_CTX_SPEED(speed) |
            SLOT_CTX_CONTEXT_ENTRIES(context_entries);

        // dev_info2: Root Hub Port Number
        // 注意：如果设备连接在root hub的port N，这里应该设置为N+1
        uint8_t root_port = device->port + 1;
        dev_priv->input_ctx_64->slot.dev_info2 =
            SLOT_CTX_ROOT_HUB_PORT(root_port);

        // tt_info: TT相关，对于高速设备为0
        dev_priv->input_ctx_64->slot.tt_info = 0;

        // dev_state: 由硬件设置
        dev_priv->input_ctx_64->slot.dev_state = 0;

        // ===== 3. 配置 EP0 上下文 =====

        // 根据速度确定最大包大小
        uint16_t max_packet_size;
        switch (speed) {
        case USB_SPEED_LOW:
            max_packet_size = 8;
            break;
        case USB_SPEED_FULL:
            max_packet_size = 8; // 初始值，后续可能更新为64
            break;
        case USB_SPEED_HIGH:
            max_packet_size = 64;
            break;
        case USB_SPEED_SUPER:
        case USB_SPEED_SUPER + 1: // Super Speed Plus
            max_packet_size = 512;
            break;
        default:
            max_packet_size = 8;
            break;
        }

        // ep_info: EP State, Interval, etc.
        // Control endpoints 没有 interval，设为0
        dev_priv->input_ctx_64->endpoints[0].ep_info = 0;

        // ep_info2: Error Count, EP Type, Max Burst, Max Packet Size
        dev_priv->input_ctx_64->endpoints[0].ep_info2 =
            EP_CTX_ERROR_COUNT(3) |           // CErr = 3
            EP_CTX_EP_TYPE(EP_TYPE_CONTROL) | // Control endpoint
            EP_CTX_MAX_BURST(0) |             // Max Burst = 0 for control
            EP_CTX_MAX_PACKET_SIZE(max_packet_size);

        // tr_dequeue_ptr: Transfer Ring地址 + DCS bit
        dev_priv->input_ctx_64->endpoints[0].tr_dequeue_ptr =
            ep0_priv->transfer_ring->phys_addr | EP_CTX_DCS;

        // tx_info: Average TRB Length
        // 对于控制端点，设置为8（setup packet大小）
        dev_priv->input_ctx_64->endpoints[0].tx_info = 8;

        device->endpoints[0].max_packet_size = max_packet_size;
    } else {
        // 清零输入上下文
        memset(dev_priv->input_ctx_32, 0, sizeof(xhci_input_ctx_t));

        // ===== 1. 设置输入控制上下文 =====
        dev_priv->input_ctx_32->ctrl.drop_flags = 0;
        dev_priv->input_ctx_32->ctrl.add_flags = 0x03; // A0 (Slot) | A1 (EP0)

        // ===== 2. 配置 Slot 上下文 =====
        uint32_t route_string = 0; // Root hub port
        uint32_t speed = usb_speed_to_xhci_port_speed(device->speed);
        uint32_t context_entries = 1; // EP0的DCI是1

        // dev_info: Route String, Speed, Context Entries
        dev_priv->input_ctx_32->slot.dev_info =
            SLOT_CTX_ROUTE_STRING(route_string) | SLOT_CTX_SPEED(speed) |
            SLOT_CTX_CONTEXT_ENTRIES(context_entries);

        // dev_info2: Root Hub Port Number
        // 注意：如果设备连接在root hub的port N，这里应该设置为N+1
        uint8_t root_port = device->port + 1;
        dev_priv->input_ctx_32->slot.dev_info2 =
            SLOT_CTX_ROOT_HUB_PORT(root_port);

        // tt_info: TT相关，对于高速设备为0
        dev_priv->input_ctx_32->slot.tt_info = 0;

        // dev_state: 由硬件设置
        dev_priv->input_ctx_32->slot.dev_state = 0;

        // ===== 3. 配置 EP0 上下文 =====

        // 根据速度确定最大包大小
        uint16_t max_packet_size;
        switch (speed) {
        case USB_SPEED_LOW:
            max_packet_size = 8;
            break;
        case USB_SPEED_FULL:
            max_packet_size = 8; // 初始值，后续可能更新为64
            break;
        case USB_SPEED_HIGH:
            max_packet_size = 64;
            break;
        case USB_SPEED_SUPER:
        case USB_SPEED_SUPER + 1: // Super Speed Plus
            max_packet_size = 512;
            break;
        default:
            max_packet_size = 8;
            break;
        }

        // ep_info: EP State, Interval, etc.
        // Control endpoints 没有 interval，设为0
        dev_priv->input_ctx_32->endpoints[0].ep_info = 0;

        // ep_info2: Error Count, EP Type, Max Burst, Max Packet Size
        dev_priv->input_ctx_32->endpoints[0].ep_info2 =
            EP_CTX_ERROR_COUNT(3) |           // CErr = 3
            EP_CTX_EP_TYPE(EP_TYPE_CONTROL) | // Control endpoint
            EP_CTX_MAX_BURST(0) |             // Max Burst = 0 for control
            EP_CTX_MAX_PACKET_SIZE(max_packet_size);

        // tr_dequeue_ptr: Transfer Ring地址 + DCS bit
        dev_priv->input_ctx_32->endpoints[0].tr_dequeue_ptr =
            ep0_priv->transfer_ring->phys_addr | EP_CTX_DCS;

        // tx_info: Average TRB Length
        // 对于控制端点，设置为8（setup packet大小）
        dev_priv->input_ctx_32->endpoints[0].tx_info = 8;

        device->endpoints[0].max_packet_size = max_packet_size;
    }

    // 验证输入上下文物理地址对齐
    if (dev_priv->input_ctx_phys & 0x3F) {
        printk("XHCI: WARNING - Input context not 64-byte aligned: 0x%lx\n",
               dev_priv->input_ctx_phys);
    }

    // ===== 4. 发送 Address Device 命令 =====

    xhci_command_completion_t *completion = xhci_alloc_command_completion();
    if (!completion) {
        printk("XHCI: Failed to allocate command completion\n");
        return -1;
    }

    completion->device = device;

    spin_lock(&xhci_command_lock);

    // Address Device TRB:
    // - Parameter: Input Context Pointer
    // - Status: Reserved (0)
    // - Control: TRB Type, Slot ID, BSR (Block Set Address Request)

    // BSR=0 表示要真正分配地址（如果BSR=1，只是启用slot但不分配地址）
    uint32_t control = (TRB_TYPE_ADDRESS_DEV << 10) |        // TRB Type
                       ((uint32_t)dev_priv->slot_id << 24) | // Slot ID
                       (0 << 9) |    // BSR = 0 (分配地址)
                       TRB_FLAG_IOC; // Interrupt on Completion

    xhci_trb_t *cmd_trb =
        xhci_queue_trb(xhci->cmd_ring,
                       dev_priv->input_ctx_phys, // Input Context pointer
                       0,                        // Status (reserved)
                       control);

    xhci_track_command(xhci, cmd_trb, completion, TRB_TYPE_ADDRESS_DEV);

    // 敲门铃
    xhci_ring_doorbell(xhci, 0, 0);

    spin_unlock(&xhci_command_lock);

    // 等待命令完成
    int ret = xhci_wait_for_command(completion, 5000);

    if (ret == 0) {
        device->address = address;
        device->state = USB_STATE_ADDRESS;
    } else {
        printk("\n*** Address Device FAILED ***\n");
        printk("Error Code: %d\n", ret);
        printk("Completion Code: %d\n", completion->completion_code);
    }

    xhci_free_command_completion(completion);
    return ret;
}

static int xhci_configure_endpoint(usb_hcd_t *hcd, usb_endpoint_t *endpoint) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;
    usb_device_t *device = endpoint->device;
    xhci_device_private_t *dev_priv =
        (xhci_device_private_t *)device->hcd_private;

    if (!dev_priv) {
        printk("XHCI: Device private data is NULL\n");
        return -1;
    }

    uint8_t ep_num = endpoint->address & 0x0F;
    uint8_t is_in = (endpoint->address & 0x80) ? 1 : 0;

    // 计算 DCI (Device Context Index)
    // EP0 的 DCI 是 1
    // 其他端点：OUT 端点 DCI = ep_num * 2, IN 端点 DCI = ep_num * 2 + 1
    uint8_t dci;
    if (ep_num == 0) {
        dci = 1; // EP0
    } else {
        dci = (ep_num * 2) + (is_in ? 1 : 0);
    }

    // 分配端点私有数据
    xhci_endpoint_private_t *ep_priv =
        (xhci_endpoint_private_t *)malloc(sizeof(xhci_endpoint_private_t));
    if (!ep_priv) {
        printk("XHCI: Failed to allocate endpoint private data\n");
        return -1;
    }

    memset(ep_priv, 0, sizeof(xhci_endpoint_private_t));
    ep_priv->dci = dci;

    // 分配传输环
    ep_priv->transfer_ring = xhci_alloc_ring(256);
    if (!ep_priv->transfer_ring) {
        printk("XHCI: Failed to allocate transfer ring\n");
        free(ep_priv);
        return -1;
    }

    if (xhci->use_64byte_context) {
        // 清空输入上下文
        memset(dev_priv->input_ctx_64, 0, sizeof(xhci_input_ctx_64_t));

        // 设置输入控制上下文
        // A0 (Slot Context) 必须设置，因为我们要更新 Context Entries
        // A[DCI] 设置我们要配置的端点
        dev_priv->input_ctx_64->ctrl.add_flags =
            (1 << dci) | (1 << 0); // A[DCI] | A0
        dev_priv->input_ctx_64->ctrl.drop_flags = 0;

        // 更新 Slot Context
        // Context Entries 必须 >= DCI（包含所有活动端点）
        uint32_t context_entries = dci; // 至少要包含这个端点

        // 保留原有的 route string 和 speed
        uint32_t route_string = dev_priv->input_ctx_64->slot.dev_info & 0xFFFFF;
        uint32_t speed = (dev_priv->input_ctx_64->slot.dev_info >> 20) & 0xF;

        // 如果这些值为0（没有初始化），使用设备的值
        if (speed == 0) {
            speed = usb_speed_to_xhci_port_speed(device->speed);
        }

        dev_priv->input_ctx_64->slot.dev_info =
            SLOT_CTX_ROUTE_STRING(route_string) | SLOT_CTX_SPEED(speed) |
            SLOT_CTX_CONTEXT_ENTRIES(context_entries);

        // 保留 Root Hub Port Number
        uint32_t root_port =
            (dev_priv->input_ctx_64->slot.dev_info2 >> 16) & 0xFF;
        if (root_port == 0) {
            root_port = device->port + 1;
        }

        dev_priv->input_ctx_64->slot.dev_info2 =
            SLOT_CTX_ROOT_HUB_PORT(root_port);

        // 配置端点上下文

        uint32_t ep_type_attr = endpoint->attributes & 0x03;
        uint32_t xhci_ep_type;

        // 确定 XHCI EP Type
        switch (ep_type_attr) {
        case USB_ENDPOINT_XFER_CONTROL:
            xhci_ep_type = EP_TYPE_CONTROL; // 4
            break;
        case USB_ENDPOINT_XFER_ISOC:
            xhci_ep_type = is_in ? EP_TYPE_ISOC_IN : EP_TYPE_ISOC_OUT; // 5 : 1
            break;
        case USB_ENDPOINT_XFER_BULK:
            xhci_ep_type = is_in ? EP_TYPE_BULK_IN : EP_TYPE_BULK_OUT; // 6 : 2
            break;
        case USB_ENDPOINT_XFER_INT:
            xhci_ep_type =
                is_in ? EP_TYPE_INTERRUPT_IN : EP_TYPE_INTERRUPT_OUT; // 7 : 3
            break;
        default:
            printk("XHCI: Unknown endpoint type: 0x%02x\n", ep_type_attr);
            xhci_free_ring(ep_priv->transfer_ring);
            free(ep_priv);
            return -1;
        }

        // 计算 Interval（对于中断和等时端点）
        uint32_t interval = 0;
        if (ep_type_attr == USB_ENDPOINT_XFER_INT ||
            ep_type_attr == USB_ENDPOINT_XFER_ISOC) {

            // 将 USB bInterval 转换为 XHCI Interval
            // XHCI Interval = log2(bInterval) for HS/SS
            // 对于 FS/LS 中断端点，转换公式不同

            uint8_t bInterval = endpoint->interval;

            if (speed == USB_SPEED_HIGH || speed == USB_SPEED_SUPER) {
                // 对于高速/超速，Interval = bInterval - 1
                if (bInterval > 0) {
                    interval = bInterval - 1;
                }
            } else {
                // 对于全速/低速，需要转换
                // bInterval 是毫秒数，需要转换为 2^n 微帧
                if (bInterval > 0) {
                    // 简化：取 log2
                    uint32_t val = bInterval;
                    interval = 0;
                    while (val > 1) {
                        val >>= 1;
                        interval++;
                    }
                    interval += 3; // 转换到微帧（8微帧 = 1毫秒）
                }
            }

            // 限制在有效范围内 (0-15)
            if (interval > 15) {
                interval = 15;
            }
        }

        // ep_info: Interval, MaxPStreams, Mult, etc.
        dev_priv->input_ctx_64->endpoints[dci - 1].ep_info =
            EP_CTX_INTERVAL(interval) |
            EP_CTX_MAX_P_STREAMS(0) | // 0 = Linear Stream Array
            EP_CTX_MULT(0);           // 0 for non-SS, not isoc

        // ep_info2: Error Count, EP Type, Max Burst Size, Max Packet Size
        dev_priv->input_ctx_64->endpoints[dci - 1].ep_info2 =
            EP_CTX_ERROR_COUNT(3) |        // CErr = 3 (retry count)
            EP_CTX_EP_TYPE(xhci_ep_type) | // EP Type
            EP_CTX_MAX_BURST(0) |          // 0 for non-SS endpoints
            EP_CTX_MAX_PACKET_SIZE(endpoint->max_packet_size);

        // tr_dequeue_ptr: Transfer Ring 物理地址 + DCS
        dev_priv->input_ctx_64->endpoints[dci - 1].tr_dequeue_ptr =
            ep_priv->transfer_ring->phys_addr | EP_CTX_DCS;

        // tx_info: Average TRB Length (用于带宽计算)
        // 对于批量端点，设置为典型值或 max packet size
        uint32_t avg_trb_length = endpoint->max_packet_size;
        if (ep_type_attr == USB_ENDPOINT_XFER_CONTROL) {
            avg_trb_length = 8; // Control setup packet
        }

        dev_priv->input_ctx_64->endpoints[dci - 1].tx_info = avg_trb_length;
    } else {
        // 清空输入上下文
        memset(dev_priv->input_ctx_32, 0, sizeof(xhci_input_ctx_t));

        // 设置输入控制上下文
        // A0 (Slot Context) 必须设置，因为我们要更新 Context Entries
        // A[DCI] 设置我们要配置的端点
        dev_priv->input_ctx_32->ctrl.add_flags =
            (1 << dci) | (1 << 0); // A[DCI] | A0
        dev_priv->input_ctx_32->ctrl.drop_flags = 0;

        // 更新 Slot Context
        // Context Entries 必须 >= DCI（包含所有活动端点）
        uint32_t context_entries = dci; // 至少要包含这个端点

        // 保留原有的 route string 和 speed
        uint32_t route_string = dev_priv->input_ctx_32->slot.dev_info & 0xFFFFF;
        uint32_t speed = (dev_priv->input_ctx_32->slot.dev_info >> 20) & 0xF;

        // 如果这些值为0（没有初始化），使用设备的值
        if (speed == 0) {
            speed = usb_speed_to_xhci_port_speed(device->speed);
        }

        dev_priv->input_ctx_32->slot.dev_info =
            SLOT_CTX_ROUTE_STRING(route_string) | SLOT_CTX_SPEED(speed) |
            SLOT_CTX_CONTEXT_ENTRIES(context_entries);

        // 保留 Root Hub Port Number
        uint32_t root_port =
            (dev_priv->input_ctx_32->slot.dev_info2 >> 16) & 0xFF;
        if (root_port == 0) {
            root_port = device->port + 1;
        }

        dev_priv->input_ctx_32->slot.dev_info2 =
            SLOT_CTX_ROOT_HUB_PORT(root_port);

        // 配置端点上下文

        uint32_t ep_type_attr = endpoint->attributes & 0x03;
        uint32_t xhci_ep_type;

        // 确定 XHCI EP Type
        switch (ep_type_attr) {
        case USB_ENDPOINT_XFER_CONTROL:
            xhci_ep_type = EP_TYPE_CONTROL; // 4
            break;
        case USB_ENDPOINT_XFER_ISOC:
            xhci_ep_type = is_in ? EP_TYPE_ISOC_IN : EP_TYPE_ISOC_OUT; // 5 : 1
            break;
        case USB_ENDPOINT_XFER_BULK:
            xhci_ep_type = is_in ? EP_TYPE_BULK_IN : EP_TYPE_BULK_OUT; // 6 : 2
            break;
        case USB_ENDPOINT_XFER_INT:
            xhci_ep_type =
                is_in ? EP_TYPE_INTERRUPT_IN : EP_TYPE_INTERRUPT_OUT; // 7 : 3
            break;
        default:
            printk("XHCI: Unknown endpoint type: 0x%02x\n", ep_type_attr);
            xhci_free_ring(ep_priv->transfer_ring);
            free(ep_priv);
            return -1;
        }

        // 计算 Interval（对于中断和等时端点）
        uint32_t interval = 0;
        if (ep_type_attr == USB_ENDPOINT_XFER_INT ||
            ep_type_attr == USB_ENDPOINT_XFER_ISOC) {

            // 将 USB bInterval 转换为 XHCI Interval
            // XHCI Interval = log2(bInterval) for HS/SS
            // 对于 FS/LS 中断端点，转换公式不同

            uint8_t bInterval = endpoint->interval;

            if (speed == USB_SPEED_HIGH || speed == USB_SPEED_SUPER) {
                // 对于高速/超速，Interval = bInterval - 1
                if (bInterval > 0) {
                    interval = bInterval - 1;
                }
            } else {
                // 对于全速/低速，需要转换
                // bInterval 是毫秒数，需要转换为 2^n 微帧
                if (bInterval > 0) {
                    // 简化：取 log2
                    uint32_t val = bInterval;
                    interval = 0;
                    while (val > 1) {
                        val >>= 1;
                        interval++;
                    }
                    interval += 3; // 转换到微帧（8微帧 = 1毫秒）
                }
            }

            // 限制在有效范围内 (0-15)
            if (interval > 15) {
                interval = 15;
            }
        }

        // ep_info: Interval, MaxPStreams, Mult, etc.
        dev_priv->input_ctx_32->endpoints[dci - 1].ep_info =
            EP_CTX_INTERVAL(interval) |
            EP_CTX_MAX_P_STREAMS(0) | // 0 = Linear Stream Array
            EP_CTX_MULT(0);           // 0 for non-SS, not isoc

        // ep_info2: Error Count, EP Type, Max Burst Size, Max Packet Size
        dev_priv->input_ctx_32->endpoints[dci - 1].ep_info2 =
            EP_CTX_ERROR_COUNT(3) |        // CErr = 3 (retry count)
            EP_CTX_EP_TYPE(xhci_ep_type) | // EP Type
            EP_CTX_MAX_BURST(0) |          // 0 for non-SS endpoints
            EP_CTX_MAX_PACKET_SIZE(endpoint->max_packet_size);

        // tr_dequeue_ptr: Transfer Ring 物理地址 + DCS
        dev_priv->input_ctx_32->endpoints[dci - 1].tr_dequeue_ptr =
            ep_priv->transfer_ring->phys_addr | EP_CTX_DCS;

        // tx_info: Average TRB Length (用于带宽计算)
        // 对于批量端点，设置为典型值或 max packet size
        uint32_t avg_trb_length = endpoint->max_packet_size;
        if (ep_type_attr == USB_ENDPOINT_XFER_CONTROL) {
            avg_trb_length = 8; // Control setup packet
        }

        dev_priv->input_ctx_32->endpoints[dci - 1].tx_info = avg_trb_length;
    }
    // 发送 Configure Endpoint 命令

    xhci_command_completion_t *completion = xhci_alloc_command_completion();
    if (!completion) {
        printk("XHCI: Failed to allocate command completion\n");
        xhci_free_ring(ep_priv->transfer_ring);
        free(ep_priv);
        return -1;
    }

    completion->device = device;

    spin_lock(&xhci_command_lock);

    // Configure Endpoint TRB
    xhci_trb_t *cmd_trb =
        xhci_queue_trb(xhci->cmd_ring, dev_priv->input_ctx_phys, 0,
                       (TRB_TYPE_CONFIG_EP << 10) |
                           ((uint32_t)dev_priv->slot_id << 24) | TRB_FLAG_IOC);

    xhci_track_command(xhci, cmd_trb, completion, TRB_TYPE_CONFIG_EP);

    xhci_ring_doorbell(xhci, 0, 0);

    spin_unlock(&xhci_command_lock);

    // 等待命令完成
    int ret = xhci_wait_for_command(completion, 5000);

    if (ret == 0) {
        endpoint->hcd_private = ep_priv;
    } else {
        printk("XHCI: Configure Endpoint failed\n");
        xhci_free_ring(ep_priv->transfer_ring);
        free(ep_priv);
    }

    xhci_free_command_completion(completion);
    return ret;
}

spinlock_t xhci_transfer_lock = {0};

// 控制传输 - 使用等待机制
static int xhci_control_transfer(usb_hcd_t *hcd, usb_transfer_t *transfer,
                                 usb_device_request_t *setup) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;
    usb_device_t *device = transfer->device;
    xhci_device_private_t *dev_priv =
        (xhci_device_private_t *)device->hcd_private;

    if (!dev_priv) {
        return -1;
    }

    uint8_t dci = 1;
    xhci_endpoint_private_t *ep_priv =
        (xhci_endpoint_private_t *)device->endpoints[0].hcd_private;

    if (!ep_priv) {
        return -1;
    }

    xhci_ring_t *ring = ep_priv->transfer_ring;

    // 分配完成结构
    xhci_transfer_completion_t *completion = xhci_alloc_transfer_completion();
    if (!completion) {
        return -1;
    }
    completion->transfer_type = NORMAL_TRANSFER;
    completion->transfer = transfer;

    spin_lock(&xhci_transfer_lock);

    // Setup Stage TRB
    uint64_t setup_data = *(uint64_t *)setup;

    xhci_trb_t *first_trb = xhci_queue_trb(
        ring, setup_data, 8,
        (TRB_TYPE_SETUP << 10) | TRB_FLAG_IDT |
            ((transfer->length > 0 && transfer->buffer)
                 ? ((setup->bmRequestType & USB_DIR_IN) ? (3 << 16) : (2 << 16))
                 : (0 << 16)));

    // Data Stage TRB
    if (transfer->length > 0 && transfer->buffer) {
        uint64_t data_addr = translate_address(get_current_page_dir(false),
                                               (uint64_t)transfer->buffer);
        uint32_t direction =
            (setup->bmRequestType & USB_DIR_IN) ? (1 << 16) : 0;
        xhci_queue_trb(ring, data_addr, transfer->length,
                       (TRB_TYPE_DATA << 10) | direction);
    }

    // Status Stage TRB
    uint32_t status_dir = (setup->bmRequestType & USB_DIR_IN) ? 0 : (1 << 16);
    if (transfer->length == 0) {
        status_dir = (setup->bmRequestType & USB_DIR_IN) ? (1 << 16) : 0;
    }

    xhci_queue_trb(ring, 0, 0,
                   (TRB_TYPE_STATUS << 10) | status_dir | TRB_FLAG_IOC);

    uint8_t first_index = first_trb - ring->trbs;

    // 跟踪传输
    xhci_track_transfer(
        xhci, first_trb,
        xhci_calc_trbs_num(ring->enqueue_index, first_index, ring->size),
        completion, transfer);

    // 敲门铃
    xhci_ring_doorbell(xhci, dev_priv->slot_id, dci);

    spin_unlock(&xhci_transfer_lock);

    // 等待传输完成（5秒超时）
    int ret = xhci_wait_for_transfer(completion, 5000);

    if (ret >= 0) {
        transfer->status = 0;
        transfer->actual_length = completion->transferred_length;
    } else {
        transfer->status = ret;
        transfer->actual_length = 0;
    }

    xhci_free_transfer_completion(completion);

    return ret;
}

// 批量传输 - 使用等待机制
static int xhci_bulk_transfer(usb_hcd_t *hcd, usb_transfer_t *transfer) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;
    usb_device_t *device = transfer->device;
    xhci_device_private_t *dev_priv =
        (xhci_device_private_t *)device->hcd_private;
    xhci_endpoint_private_t *ep_priv =
        (xhci_endpoint_private_t *)transfer->endpoint->hcd_private;

    if (!dev_priv || !ep_priv) {
        return -1;
    }

    xhci_ring_t *ring = ep_priv->transfer_ring;

    // 分配完成结构
    xhci_transfer_completion_t *completion = xhci_alloc_transfer_completion();
    if (!completion) {
        return -1;
    }
    completion->transfer_type = NORMAL_TRANSFER;
    completion->transfer = transfer;

    spin_lock(&xhci_transfer_lock);

    // 获取第一个TRB指针
    xhci_trb_t *first_trb = NULL;

    uint16_t max_packet_size = transfer->endpoint->max_packet_size;

    size_t td_packet_count =
        (transfer->length + max_packet_size - 1) / max_packet_size;

    size_t progress = 0;

    while (progress < transfer->length) {
        uint64_t data_addr = translate_address(
            get_current_page_dir(false), (uint64_t)transfer->buffer + progress);

        size_t chunk = MIN(transfer->length - progress,
                           DEFAULT_PAGE_SIZE - (progress % DEFAULT_PAGE_SIZE));

        bool chain = (progress + chunk) < transfer->length;

        size_t td_size = td_packet_count - (progress + chunk) / max_packet_size;
        if ((progress + chunk) >= transfer->length)
            td_size = 0;

        // Normal TRB
        if (td_size > 31)
            td_size = 31;
        xhci_trb_t *trb = xhci_queue_trb(
            ring, data_addr, chunk | (td_size << 17),
            (TRB_TYPE_NORMAL << 10) | (chain ? TRB_FLAG_CHAIN : TRB_FLAG_IOC));

        if (!first_trb) {
            first_trb = trb;
        }

        progress += chunk;
    }

    uint8_t first_index = first_trb - ring->trbs;

    // 跟踪传输
    xhci_track_transfer(
        xhci, first_trb,
        xhci_calc_trbs_num(ring->enqueue_index, first_index, ring->size),
        completion, transfer);

    // 敲门铃
    xhci_ring_doorbell(xhci, dev_priv->slot_id, ep_priv->dci);

    spin_unlock(&xhci_transfer_lock);

    // 等待传输完成
    int ret = xhci_wait_for_transfer(completion, 5000);

    if (ret >= 0) {
        transfer->status = 0;
        transfer->actual_length = completion->transferred_length;
    } else {
        transfer->status = ret;
        transfer->actual_length = 0;
    }

    xhci_free_transfer_completion(completion);

    return ret;
}

// 中断传输 - 使用需放到单独线程
static int xhci_interrupt_transfer(usb_hcd_t *hcd, usb_transfer_t *transfer) {
    xhci_hcd_t *xhci = (xhci_hcd_t *)hcd->private_data;
    usb_device_t *device = transfer->device;
    xhci_device_private_t *dev_priv =
        (xhci_device_private_t *)device->hcd_private;
    xhci_endpoint_private_t *ep_priv =
        (xhci_endpoint_private_t *)transfer->endpoint->hcd_private;

    if (!dev_priv || !ep_priv) {
        return -1;
    }

    xhci_ring_t *ring = ep_priv->transfer_ring;

    // 分配完成结构
    xhci_transfer_completion_t *completion = xhci_alloc_transfer_completion();
    if (!completion) {
        return -1;
    }
    completion->transfer_type = INTR_TRANSFER;
    completion->transfer = transfer;

    spin_lock(&xhci_transfer_lock);

    // Normal TRB
    uint64_t data_addr = translate_address(get_current_page_dir(false),
                                           (uint64_t)transfer->buffer);

    uint32_t control = (TRB_TYPE_NORMAL << 10) | TRB_FLAG_IOC;

    xhci_trb_t *first_trb =
        xhci_queue_trb(ring, data_addr, transfer->length, control);

    uint8_t first_index = first_trb - ring->trbs;

    // 跟踪传输
    xhci_track_transfer(
        xhci, first_trb,
        xhci_calc_trbs_num(ring->enqueue_index, first_index, ring->size),
        completion, transfer);

    // 敲门铃
    xhci_ring_doorbell(xhci, dev_priv->slot_id, ep_priv->dci);

    spin_unlock(&xhci_transfer_lock);

    int ret = xhci_wait_for_transfer(completion, 0);

    if (ret >= 0) {
        transfer->status = 0;
        transfer->actual_length = completion->transferred_length;
    } else {
        transfer->status = ret;
        transfer->actual_length = 0;
    }

    xhci_free_transfer_completion(completion);

    return 0;
}

typedef struct enumerater_arg {
    usb_hcd_t *hcd;
    uint8_t port_id;
} enumerater_arg_t;

spinlock_t enumerate_lock = {0};

void xhci_device_enumerater(enumerater_arg_t *arg) {
    spin_lock(&enumerate_lock);

    usb_hcd_t *hcd = arg->hcd;

    xhci_hcd_t *xhci = hcd->private_data;

    if (hcd) {
        int ret = usb_enumerate_device(hcd, xhci->root_hub, arg->port_id);
        if (ret == 1) {
            printk("XHCI: Device on port %d enumerated successfully\n",
                   arg->port_id);
        } else {
            printk("XHCI: Failed to enumerate device on port %d\n",
                   arg->port_id);
        }
    }

    free(arg);

    spin_unlock(&enumerate_lock);

    task_exit(0);
}

// 处理端口状态变化
void xhci_handle_port_status(xhci_hcd_t *xhci, uint8_t port_id) {
    if (port_id >= xhci->max_ports) {
        return;
    }

    uint32_t portsc = xhci_readl(&xhci->port_regs[port_id].portsc);

    // printk("XHCI: Port %d status: 0x%08x\n", port_id, portsc);

    // 连接状态变化
    if ((portsc & XHCI_PORTSC_CCS) && !xhci->connection[port_id]) {
        // 设备连接
        printk("XHCI: Device connected on port %d\n", port_id);

        // 等待端口稳定
        uint64_t target_time = 100000000ULL + nanoTime(); // 100ms
        while (nanoTime() < target_time) {
            arch_yield();
        }

        xhci->connection[port_id] = true;

        enumerater_arg_t *arg = malloc(sizeof(enumerater_arg_t));
        arg->hcd = xhci->hcd;
        arg->port_id = port_id;

        task_create("enumerater", (void (*)(uint64_t))xhci_device_enumerater,
                    (uint64_t)arg, KTHREAD_PRIORITY);
    } else if (!(portsc & XHCI_PORTSC_CCS) && xhci->connection[port_id]) {
        printk("XHCI: Device disconnected on port %d\n", port_id);
        xhci->connection[port_id] = false;
    }

    if (portsc & XHCI_PORTSC_PEC) {
        // 端口使能变化
        printk("XHCI: Port %d enable changed\n", port_id);
        xhci_writel(&xhci->port_regs[port_id].portsc, portsc | XHCI_PORTSC_PEC);
    }

    // if (portsc & XHCI_PORTSC_PRC) {
    //     // 端口重置完成
    //     printk("XHCI: Port %d reset complete\n", port_id);
    //     xhci_writel(&xhci->port_regs[port_id].portsc, portsc |
    //     XHCI_PORTSC_PRC);
    // }
}

// 初始化XHCI驱动
usb_hcd_t *xhci_init(void *mmio_base, pci_device_t *pci_dev) {
    // 分配XHCI上下文
    xhci_hcd_t *xhci = (xhci_hcd_t *)malloc(sizeof(xhci_hcd_t));
    if (!xhci) {
        printk("XHCI: Failed to allocate context\n");
        return NULL;
    }
    memset(xhci, 0, sizeof(xhci_hcd_t));

    // 初始化寄存器指针
    xhci->cap_regs = (xhci_cap_regs_t *)mmio_base;

    uint8_t caplength = xhci->cap_regs->caplength;
    uint32_t rtsoff = xhci_readl(&xhci->cap_regs->rtsoff);
    uint32_t dboff = xhci_readl(&xhci->cap_regs->dboff);

    xhci->op_regs = (xhci_op_regs_t *)((uint8_t *)mmio_base + caplength);
    xhci->runtime_regs = (xhci_runtime_regs_t *)((uint8_t *)mmio_base + rtsoff);
    xhci->doorbell_regs = (uint32_t *)((uint8_t *)mmio_base + dboff);

    xhci->port_regs = (xhci_port_regs_t *)((uint8_t *)xhci->op_regs + 0x400);
    xhci->intr_regs =
        (xhci_intr_regs_t *)((uint8_t *)xhci->runtime_regs + 0x20);

    memset(xhci->connection, 0, sizeof(xhci->connection));

    xhci->pci_dev = pci_dev;

    // 注册HCD
    usb_hcd_t *hcd = usb_register_hcd("xhci", &xhci_ops, mmio_base, xhci);
    if (!hcd) {
        printk("XHCI: Failed to register HCD\n");
        free(xhci);
        return NULL;
    }

    return hcd;
}

// 关闭XHCI驱动
void xhci_shutdown(usb_hcd_t *hcd) {
    if (hcd) {
        usb_unregister_hcd(hcd);
    }
}

int xhci_hcd_driver_probe(pci_device_t *pci_dev, uint32_t vendor_device_id) {
    // Find MMIO BAR
    uint64_t mmio_base = 0;
    uint64_t mmio_size = 0;
    for (int i = 0; i < 6; i++) {
        if (pci_dev->bars[i].size > 0 && pci_dev->bars[i].mmio) {
            mmio_base = pci_dev->bars[i].address;
            mmio_size = pci_dev->bars[i].size;
            break;
        }
    }

    if (mmio_base == 0) {
        printk("xhci: No MMIO BAR found\n");
        return -1;
    }

    // Map MMIO region
    void *mmio_vaddr = (void *)phys_to_virt(mmio_base);
    map_page_range(
        get_current_page_dir(false), (uint64_t)mmio_vaddr, mmio_base, mmio_size,
        PT_FLAG_R | PT_FLAG_W | PT_FLAG_UNCACHEABLE | PT_FLAG_DEVICE);

    usb_hcd_t *xhci_hcd = xhci_init(mmio_vaddr, pci_dev);

    return 0;
}

void xhci_hcd_driver_remove(pci_device_t *dev) { xhci_shutdown(dev->desc); }

void xhci_hcd_driver_shutdown(pci_device_t *dev) { xhci_shutdown(dev->desc); }

pci_driver_t xhci_hcd_driver = {
    .name = "xhci_hcd",
    .class_id = 0x000C0330,
    .vendor_device_id = 0x00000000,
    .probe = xhci_hcd_driver_probe,
    .remove = xhci_hcd_driver_remove,
    .shutdown = xhci_hcd_driver_shutdown,
    .flags = 0,
};

__attribute__((visibility("default"))) int dlmain() {
    regist_pci_driver(&xhci_hcd_driver);

    return 0;
}
